#!/usr/bin/python
# Aaron Saunders | ams@bio.aau.dk

import sys
from Bio import SeqIO
from Bio.SeqUtils.CheckSum import seguid
import matplotlib.pyplot as plt
import pickle

infname = sys.argv[1]
outfilestub = sys.argv[2]

reps = {}
counts = {}
uniquebuffer = []
uniquefname = outfilestub + '.unique.fastq'
uniquefile = open(uniquefname, 'w')

for seqrecord in SeqIO.parse(open(infname, 'r'), 'fasta'):
    checksum = seguid(seqrecord.seq)
    if checksum not in reps:
        reps[checksum] = seqrecord.id
        uniquebuffer.append(seqrecord)
        counts[checksum] = 1
    else:
        counts[checksum] += 1
    if len(reps) % 1000 == 0:
        SeqIO.write(uniquebuffer, uniquefile, 'fasta')

SeqIO.write(uniquebuffer, uniquefile, 'fasta')
uniquefile.close()

with open(outfilestub + 'counts.pkl', 'wb') as fh:
    pickle.dump(counts, fh)

plt.hist(counts.values(), bins=10, align='mid')
plt.xlabel('counts')
plt.ylabel('#seqs')
plt.savefig(outfilestub + '_histo.png', format='png')
